package generator

import (
	"kumachan/interpreter/lang/common/source"
	"kumachan/interpreter/compiler/checker/checked"
	rt "kumachan/interpreter/runtime/def"
)


type ExprContext struct {
	Code       ExprCode
	Scope      LocalScope
	Static     StaticCollector
	GenStatic  GeneratedStaticCollector
	*Statistics
}
func makeExprContext (
	input checked.ProductPatternInfo,
	inj   [] *checked.LocalBinding,
	self  *checked.LocalBinding,
	s     *Statistics,
) *ExprContext {
	var ctx ExprContext
	var args = make([] ArgumentInfo, len(input))
	for i, item := range input {
		args[i] = ArgumentInfo {
			Binding: item.Binding,
			Index1:  item.Index1,
			Used:    false,
		}
	}
	ctx.Scope.Arguments = args
	ctx.Scope.Inject = inj
	ctx.Scope.SelfRef = self
	ctx.Statistics = s
	return &ctx
}
func MakeFunctionExprContext(input checked.ProductPatternInfo, inj ([] *checked.LocalBinding), s *Statistics) *ExprContext {
	return makeExprContext(input, inj, nil, s)
}
func MakeLambdaExprContext(parent *ExprContext, input checked.ProductPatternInfo, self *checked.LocalBinding) *ExprContext {
	return makeExprContext(input, nil, self, parent.Statistics)
}
func (ctx *ExprContext) CollectEverything(info rt.FunctionInfo) (
	*rt.FunctionSeedUsual,
	([] *checked.LocalBinding), // captured bindings
) {
	var args, captures = ctx.Scope.Collect(ctx.Statistics)
	var capture_quantity = uint(len(captures))
	var injection_quantity = uint(len(ctx.Scope.Inject))
	var ctx_length = rt.MustBeLocalAddr(capture_quantity + injection_quantity)
	var static = ctx.Static.Collect()
	var gen_static = ctx.GenStatic.Collect()
	var seed = &rt.FunctionSeedUsual {
		InstTree:  ctx.Code.InstTree,
		Position:  ctx.Code.Position,
		ExtBranch: ctx.Code.ExtBranch,
		ExtUpdate: ctx.Code.ExtUpdate,
		Arguments: args,
		CtxSize:   ctx_length,
		Lambdas:   ctx.Code.Lambdas,
		Static:    static,
		GenStatic: gen_static,
		Info:      info,
	}
	return seed, captures
}

type LocalScope struct {
	Captures   map[*checked.LocalBinding] uint
	Arguments  [] ArgumentInfo
	Inject     [] *checked.LocalBinding
	SelfRef    *checked.LocalBinding  // nullable
}
type ArgumentInfo struct {
	Binding  *checked.LocalBinding
	Index1   uint
	Used     bool
}
func (scope *LocalScope) HasInject() bool {
	return (len(scope.Inject) > 0)
}
func (scope *LocalScope) HasSelfRef() bool {
	return (scope.SelfRef != nil)
}
func (scope *LocalScope) GetArgIndex(b *checked.LocalBinding) (uint, bool) {
	for i := range scope.Arguments {
		var arg = &(scope.Arguments[i])
		if b == arg.Binding {
			arg.Used = true
			return uint(i), true
		}
	}
	var shift = 0
	if scope.HasSelfRef() {
		if b == scope.SelfRef {
			return 0, false
		} else {
			shift += 1
		}
	}
	if scope.HasInject() {
		if shift > 0 {
			// combined generated context (self + inject) not needed now
			panic("something went wrong")
		}
		for i, item := range scope.Inject {
			if b == item {
				return uint(shift + i), false
			}
		}
		shift += len(scope.Inject)
	}
	if scope.Captures == nil {
		scope.Captures = make(map[*checked.LocalBinding] uint)
	}
	var index, exists = scope.Captures[b]
	if exists {
		return index, false
	} else {
		var index = uint(len(scope.Captures) + shift)
		scope.Captures[b] = index
		return index, false
	}
}
func (scope *LocalScope) Collect(s *Statistics) ([] rt.ShortIndex, [] *checked.LocalBinding) {
	var args = make([] rt.ShortIndex, len(scope.Arguments))
	var captures = make([] *checked.LocalBinding, len(scope.Captures))
	for i, item := range scope.Arguments {
		if !(item.Used) {
			if item.Binding != nil {
				s.MarkUnused(item.Binding)
			}
		}
		args[i] = rt.EnsureShortIndex(item.Index1)
	}
	var occurred = make(map[uint] struct{})
	for item, i := range scope.Captures {
		var _, duplicate = occurred[i]
		if duplicate { panic("something went wrong") }
		occurred[i] = struct{}{}
		var index = (func() uint {
			var index = i
			if scope.HasSelfRef() {
				index -= 1
			}
			if scope.HasInject() {
				index -= uint(len(scope.Inject))
			}
			return index
		})()
		captures[index] = item
	}
	return args, captures
}

type StaticCollector struct {
	List   [] rt.StaticValueSeed
	Cache  map[string] rt.LocalAddr
}
func (sc *StaticCollector) Register(seed rt.StaticValueSeed) rt.LocalAddr {
	var key, has_key = seed.Key()
	if has_key {
		var cached_addr, exists = sc.Cache[key]
		if exists {
			return cached_addr
		}
	}
	sc.List = append(sc.List, seed)
	var addr = rt.MustBeLocalAddr(uint(len(sc.List) - 1))
	if has_key {
		if sc.Cache == nil { sc.Cache = make(map[string] rt.LocalAddr) }
		sc.Cache[key] = addr
	}
	return addr
}
func (sc *StaticCollector) Collect() ([] rt.StaticValueSeed) {
	return sc.List
}

type GeneratedStaticCollector struct {
	PathList    [] rt.InterfaceTransformPath
	TableList   [] rt.DispatchTableSeed
	TableCache  map[rt.DispatchTableSeed] rt.LocalAddr
}
func (gsc *GeneratedStaticCollector) RegisterInterfaceTransformPath(path checked.InterfaceTransformPath) rt.LocalAddr {
	gsc.PathList = append(gsc.PathList, ([] uint)(path))
	return rt.MustBeLocalAddr(uint(len(gsc.PathList)))
}
func (gsc *GeneratedStaticCollector) RegisterDispatchTable(d *checked.DispatchTable) rt.LocalAddr {
	var id = rt.DispatchTableIdentifier {
		ConcreteTypeId:  typeNameToSymbol(d.Concrete),
		InterfaceTypeId: typeNameToSymbol(d.Interface),
	}
	var seed = rt.DispatchTableSeed { Id: id }
	var cached_addr, exists = gsc.TableCache[seed]
	if exists {
		return cached_addr
	} else {
		gsc.TableList = append(gsc.TableList, seed)
		var addr = rt.MustBeLocalAddr(uint(len(gsc.TableList) - 1))
		if gsc.TableCache == nil { gsc.TableCache = make(map[rt.DispatchTableSeed] rt.LocalAddr) }
		gsc.TableCache[seed] = addr
		return addr
	}
}
func (gsc *GeneratedStaticCollector) Collect() rt.GeneratedStaticDataSeed {
	return rt.GeneratedStaticDataSeed {
		DispatchTables: gsc.TableList,
		TransformPaths: gsc.PathList,
	}
}

type ExprCode struct {
	InstTree   [] rt.Instruction
	Position   [] source.Position
	ExtBranch  [] rt.ExternalBranchMapping
	ExtUpdate  [] rt.ExternalUpdateIndexes
	Lambdas    [] *rt.FunctionSeedUsual
}
func (code *ExprCode) Write(pos source.Position, inst rt.Instruction) rt.LocalAddr {
	code.InstTree = append(code.InstTree, inst)
	code.Position = append(code.Position, pos)
	return rt.MustBeLocalAddr(uint(len(code.InstTree) - 1))
}
func (code *ExprCode) WriteVec(pos source.Position, elements ([] rt.LocalAddr)) rt.LocalAddr {
	var length = rt.MustBeLocalAddr(uint(len(elements)))
	var length_addr = code.Write(pos, rt.MakeSizeInstruction(length))
	for _, element_addr := range elements {
		code.Write(pos, rt.Instruction {
			OpCode: rt.FWD,
			Obj:    element_addr,
		})
	}
	return length_addr
}
func (code *ExprCode) AppendBranchMapping(mapping rt.ExternalBranchMapping) rt.ExternalCodeEntityIndex {
	code.ExtBranch = append(code.ExtBranch, mapping)
	return rt.RequireExternalCodeEntityIndex(uint(len(code.ExtBranch) - 1))
}
func (code *ExprCode) AppendUpdateIndexes(indexes rt.ExternalUpdateIndexes) rt.ExternalCodeEntityIndex {
	code.ExtUpdate = append(code.ExtUpdate, indexes)
	return rt.RequireExternalCodeEntityIndex(uint(len(code.ExtUpdate) - 1))
}
func (code *ExprCode) AddLambda(seed *rt.FunctionSeedUsual) rt.LocalAddr {
	code.Lambdas = append(code.Lambdas, seed)
	return rt.MustBeLocalAddr(uint(len(code.Lambdas) - 1))
}

func compileLambda(f checked.Lambda, loc source.Location, ctx *ExprContext) (
	rt.LocalAddr, ([] *checked.LocalBinding),
) {
	var info = rt.FunctionInfo { Location: loc }
	var inner_ctx = MakeLambdaExprContext(ctx, f.In, f.Self)
	compileExpr(f.Out, inner_ctx)
	var seed, captures = inner_ctx.CollectEverything(info)
	return ctx.Code.AddLambda(seed), captures
}

func compileExpr(expr *checked.Expr, ctx *ExprContext) rt.LocalAddr {
	var pos = expr.Info.Location.Pos
	var use = func(inst rt.Instruction) rt.LocalAddr {
		var addr = ctx.Code.Write(pos, inst)
		return addr
	}
	var compile = func(child *checked.Expr) rt.LocalAddr {
		return compileExpr(child, ctx)
	}
	var compileAll = func(children ([] *checked.Expr)) rt.LocalAddr {
		var addr_vec = make([] rt.LocalAddr, len(children))
		for i, child := range children {
			addr_vec[i] = compile(child)
		}
		return ctx.Code.WriteVec(pos, addr_vec)
	}
	var appendBranchMapping = func(mapping rt.ExternalBranchMapping) rt.ExternalCodeEntityIndex {
		return ctx.Code.AppendBranchMapping(mapping)
	}
	var appendUpdateIndexes = func(indexes rt.ExternalUpdateIndexes) rt.ExternalCodeEntityIndex {
		return ctx.Code.AppendUpdateIndexes(indexes)
	}
	switch content := expr.Content.(type) {
	case checked.Argument:
		return use(rt.Instruction {
			OpCode: rt.ARG,
			Idx:    0,
		})
	case checked.Block:
		var nested = checked.NestedLambdaCalls(content, expr.Info)
		return compile(nested)
	case checked.Branch:
		var lambda_expr = &checked.Expr {
			Type:    nil,
			Info:    expr.Info,
			Content: content.Lambda,
		}
		return compile(lambda_expr)
	case checked.Call:
		return use(rt.Instruction {
			OpCode: rt.CALL,
			Obj:    compile(content.Callee),
			Src:    compile(content.Argument),
		})
	case checked.Enum:
		return use(rt.Instruction {
			OpCode: rt.ENUM,
			Idx:    rt.EnsureShortIndex(content.Index),
			Obj:    compile(content.Value),
		})
	case checked.FuncRef:
		var sym = functionNameToSymbol(content.Name)
		ctx.Statistics.MarkDepTo(sym)
		var seed = rt.StaticValueSeedFunctionReference { Symbol: sym }
		var static = rt.Instruction {
			OpCode: rt.STATIC,
			Src:    ctx.Static.Register(seed),
		}
		if len(content.Implicit) > 0 {
			return use(rt.Instruction {
				OpCode: rt.INJ,
				Obj:    compileAll(content.Implicit),
				Src:    use(static),
			})
		} else {
			return use(static)
		}
	case checked.Interface:
		return use(rt.Instruction {
			OpCode: rt.CI,
			Obj:    compile(content.ConcreteValue),
			Src:    ctx.GenStatic.RegisterDispatchTable(content.DispatchTable),
		})
	case checked.InterfaceMethod:
		return use(rt.Instruction {
			OpCode: rt.MET,
			Idx:    rt.EnsureShortIndex(content.Index),
			Obj:    compile(content.Interface),
		})
	case checked.InterfaceTransform:
		return use(rt.Instruction {
			OpCode: rt.CII,
			Obj:    compile(content.Argument),
			Src:    ctx.GenStatic.RegisterInterfaceTransformPath(content.Path),
		})
	case checked.InteriorRef:
		var op = (func() rt.OpCode {
			switch content.Kind {
			case checked.RK_Proj:
				switch content.Operand {
				case checked.RO_Record:  return rt.PR
				case checked.RO_ProjRef: return rt.PRP
				}
			case checked.RK_Case:
				switch content.Operand {
				case checked.RO_Enum:    return rt.CR
				case checked.RO_ProjRef: return rt.CRP
				case checked.RO_CaseRef: return rt.CRC
				}
			}
			panic("invalid interior ref")
		})()
		return use(rt.Instruction {
			OpCode: op,
			Idx:    rt.EnsureShortIndex(content.Index),
			Obj:    compile(content.Base),
		})
	case checked.Lambda:
		var op = (func() rt.OpCode {
			if content.Self != nil {
				return rt.CLR
			} else {
				return rt.CL
			}
		})()
		var loc = expr.Info.Location
		var lambda, captures = compileLambda(content, loc, ctx)
		var capture_expr_list = make([] *checked.Expr, len(captures))
		for i, item := range captures {
			capture_expr_list[i] = &checked.Expr {
				Type:    nil,
				Info:    expr.Info,
				Content: checked.LocalRef { Binding: item },
			}
		}
		var context = compileAll(capture_expr_list)
		return use(rt.Instruction {
			OpCode: op,
			Obj:    context,
			Src:    lambda,
		})
	case checked.List:
		var compact_item_t, compact = (func() (rt.ShortIndex, bool) {
			if content.Info.Compact {
				var item_t = content.Info.ItemType
				var t, ok = rt.GetCompactArrayItemTypeIndex(item_t)
				return t, ok
			} else {
				return 0, false
			}
		})()
		if compact {
			return use(rt.Instruction {
				OpCode: rt.LSC,
				Idx:    compact_item_t,
				Obj:    compileAll(content.Items),
			})
		} else {
			return use(rt.Instruction {
				OpCode: rt.LSV,
				Obj:    compileAll(content.Items),
			})
		}
	case checked.LocalRef:
		var index, is_arg = ctx.Scope.GetArgIndex(content.Binding)
		if is_arg {
			return use(rt.Instruction {
				OpCode: rt.ARG,
				Idx:    rt.EnsureShortIndex(index),
			})
		} else {
			return use(rt.Instruction {
				OpCode: rt.CTX,
				Src:    rt.MustBeLocalAddr(index),
			})
		}
	case checked.NumericLiteral:
		var seed = rt.StaticValueSeedVariant {
			Value: content.Value,
		}
		return use(rt.Instruction {
			OpCode: rt.STATIC,
			Src:    ctx.Static.Register(seed),
		})
	case checked.Select:
		var mapping = (func() rt.ExternalBranchMapping {
			var arity = len(content.Arguments)
			var raw_keys = content.SortedBranchKeys()
			var m = make(map[rt.BranchKey] rt.LocalAddr)
			for _, raw_key := range raw_keys {
				var branch = content.Branches[raw_key]
				var branch_expr = branch.ToExpr()
				var branch_addr = compile(branch_expr)
				var short_key = make([] rt.ShortIndex, arity)
				for i := 0; i < arity; i += 1 {
					short_key[i] = rt.EnsureShortIndex(raw_key[i])
				}
				var key = rt.MakeBranchKey(short_key)
				var _, duplicate = m[key]
				if duplicate { panic("something went wrong") }
				m[key] = branch_addr
			}
			if content.HasDefault {
				return rt.ExternalBranchMapping {
					HasDefault: true,
					Default:    compile(content.Default.ToExpr()),
					AddrMap:    m,
				}
			} else {
				return rt.ExternalBranchMapping {
					AddrMap: m,
				}
			}
		})()
		return use(rt.Instruction {
			OpCode: rt.SELECT,
			ExtIdx: appendBranchMapping(mapping),
			Obj:    compileAll(content.Arguments),
		})
	case checked.StringFormatter:
		var seed = rt.StaticValueSeedStringFormatter {
			Template:   content.Template,
			Parameters: content.Parameters,
		}
		return use(rt.Instruction {
			OpCode: rt.STATIC,
			Src:    ctx.Static.Register(seed),
		})
	case checked.StringLiteral:
		var seed = rt.StaticValueSeedStringLiteral {
			Value: content.Value,
		}
		return use(rt.Instruction {
			OpCode: rt.STATIC,
			Src:    ctx.Static.Register(seed),
		})
	case checked.Switch:
		var mapping = (func() rt.ExternalBranchMapping {
			var raw_keys = content.SortedBranchKeys()
			var m = make(map[rt.BranchKey] rt.LocalAddr)
			for _, raw_key := range raw_keys {
				var branch = content.Branches[raw_key]
				var branch_expr = branch.ToExpr()
				var branch_addr = compile(branch_expr)
				var short_key = rt.EnsureShortIndex(raw_key)
				var key = rt.MakeBranchKeySingleItem(short_key)
				var _, duplicate = m[key]
				if duplicate { panic("something went wrong") }
				m[key] = branch_addr
			}
			if content.HasDefault {
				return rt.ExternalBranchMapping {
					HasDefault: true,
					Default:    compile(content.Default.ToExpr()),
					AddrMap:    m,
				}
			} else {
				return rt.ExternalBranchMapping {
					AddrMap: m,
				}
			}
		})()
		return use(rt.Instruction {
			OpCode: rt.SWITCH,
			ExtIdx: appendBranchMapping(mapping),
			Obj:    compile(content.Argument),
		})
	case checked.Tuple:
		var elements = content.Elements
		var parallel = content.Parallel
		var op = (func() rt.OpCode {
			if parallel {
				return rt.TVP
			} else {
				return rt.TV
			}
		})()
		return use(rt.Instruction {
			OpCode: op,
			Obj:    compileAll(elements),
		})
	case checked.TupleGet:
		return use(rt.Instruction {
			OpCode: rt.GET,
			Idx:    rt.EnsureShortIndex(content.Index),
			Obj:    compile(content.Tuple),
		})
	case checked.TupleUpdate:
		var base = content.Tuple
		var L = len(content.Replaced)
		var indexes = make(rt.ExternalUpdateIndexes, L)
		var values = make([] *checked.Expr, L)
		for i, replaced := range content.Replaced {
			indexes[i] = rt.EnsureShortIndex(replaced.Index)
			values[i] = replaced.Value
		}
		return use(rt.Instruction {
			OpCode: rt.SET,
			ExtIdx: appendUpdateIndexes(indexes),
			Obj:    compile(base),
			Src:    compileAll(values),
		})
	case checked.UnitValue:
		return use(rt.Instruction {
			OpCode: rt.NIL,
		})
	case checked.Wrapper:
		return compile(content.Inner)
	default:
		panic("impossible branch")
	}
}


